use arrow::array::{Array, ArrayRef, BooleanArray, Date64Array, Float32Array, Float64Array, Int32Array, Int64Array, StringArray, TimestampNanosecondArray};
use arrow::datatypes::{Schema, Field, DataType as ArrowDataType};
use arrow::record_batch::RecordBatch;
use chrono::{DateTime, Utc, TimeZone};
use std::sync::Arc;
use anyhow::{Result, anyhow};
use serde::{Serialize, de::DeserializeOwned};
use crate::bridge::arrow::types::{DataType, BridgeError};
use pyo3::prelude::*;
use pyo3::types::{PyDict, PyList};
use pyo3::conversion::{IntoPy, ToPyObject};
use pyo3::PyResult;

/// 基础数据类型转换器
pub trait ArrowConverter<T> {
    /// 将Rust类型转换为Arrow数组
    fn to_arrow_array(data: &[T]) -> Result<ArrayRef>;
    
    /// 将Arrow数组转换为Rust类型
    fn from_arrow_array(array: &ArrayRef) -> Result<Vec<T>>;
    
    /// 获取Arrow数据类型
    fn get_arrow_type() -> ArrowDataType;
}

/// 实现基础数据类型的Arrow转换
macro_rules! impl_arrow_converter {
    ($rust_type:ty, $arrow_type:ty, $arrow_data_type:expr, $array_type:ty, $conversion:expr) => {
        impl ArrowConverter<$rust_type> for $arrow_type {
            fn to_arrow_array(data: &[$rust_type]) -> Result<ArrayRef> {
                let array = <$array_type>::from(($conversion)(data));
                Ok(Arc::new(array))
            }
            
            fn from_arrow_array(array: &ArrayRef) -> Result<Vec<$rust_type>> {
                let typed_array = array.as_any().downcast_ref::<$array_type>()
                    .ok_or_else(|| anyhow!("Failed to downcast array"))?;
                    
                Ok((0..typed_array.len())
                    .map(|i| typed_array.value(i))
                    .collect())
            }
            
            fn get_arrow_type() -> ArrowDataType {
                $arrow_data_type
            }
        }
    };
}

// 实现基础数据类型转换
impl_arrow_converter!(bool, BooleanArray, ArrowDataType::Boolean, BooleanArray, 
    |v: &[bool]| v.to_vec());
impl_arrow_converter!(i32, Int32Array, ArrowDataType::Int32, Int32Array,
    |v: &[i32]| v.to_vec());
impl_arrow_converter!(i64, Int64Array, ArrowDataType::Int64, Int64Array,
    |v: &[i64]| v.to_vec());
impl_arrow_converter!(f32, Float32Array, ArrowDataType::Float32, Float32Array,
    |v: &[f32]| v.to_vec());
impl_arrow_converter!(f64, Float64Array, ArrowDataType::Float64, Float64Array,
    |v: &[f64]| v.to_vec());

// 字符串类型的特殊实现
impl ArrowConverter<String> for StringArray {
    fn to_arrow_array(data: &[String]) -> Result<ArrayRef> {
        let array = StringArray::from(data.iter().map(|s| s.as_str()).collect::<Vec<&str>>());
        Ok(Arc::new(array))
    }
    
    fn from_arrow_array(array: &ArrayRef) -> Result<Vec<String>> {
        let string_array = array.as_any().downcast_ref::<StringArray>()
            .ok_or_else(|| anyhow!("Failed to downcast string array"))?;
            
        Ok((0..string_array.len())
            .map(|i| string_array.value(i).to_string())
            .collect())
    }
    
    fn get_arrow_type() -> ArrowDataType {
        ArrowDataType::Utf8
    }
}

// DateTime类型的特殊实现
impl ArrowConverter<DateTime<Utc>> for TimestampNanosecondArray {
    fn to_arrow_array(data: &[DateTime<Utc>]) -> Result<ArrayRef> {
        let nanos: Vec<i64> = data.iter()
            .map(|dt| dt.timestamp_nanos_opt().unwrap_or(0))
            .collect();
        Ok(Arc::new(TimestampNanosecondArray::from(nanos)))
    }
    
    fn from_arrow_array(array: &ArrayRef) -> Result<Vec<DateTime<Utc>>> {
        let ts_array = array.as_any().downcast_ref::<TimestampNanosecondArray>()
            .ok_or_else(|| anyhow!("Failed to downcast timestamp array"))?;
            
        Ok((0..ts_array.len())
            .map(|i| Utc.timestamp_nanos(ts_array.value(i)))
            .collect())
    }
    
    fn get_arrow_type() -> ArrowDataType {
        ArrowDataType::Timestamp(arrow::datatypes::TimeUnit::Nanosecond, None)
    }
}

/// 通用结构体序列化器
pub struct StructSerializer {
    schema: Arc<Schema>,
}

impl StructSerializer {
    /// 创建新的结构体序列化器
    pub fn new(fields: Vec<Field>) -> Self {
        Self {
            schema: Arc::new(Schema::new(fields)),
        }
    }
    
    /// 序列化结构体为RecordBatch
    pub fn serialize<T: Serialize>(&self, data: &T) -> Result<RecordBatch> {
        let json = serde_json::to_value(data)?;
        let arrays = self.schema.fields().iter()
            .map(|field| {
                let field_name = field.name();
                let field_value = &json[field_name];
                self.value_to_array(field_value, field.data_type())
            })
            .collect::<Result<Vec<_>>>()?;
            
        RecordBatch::try_new(self.schema.clone(), arrays)
            .map_err(|e| anyhow!("Failed to create record batch: {}", e))
    }
    
    /// 反序列化RecordBatch为结构体
    pub fn deserialize<T: DeserializeOwned>(&self, batch: &RecordBatch) -> Result<T> {
        let mut map = serde_json::Map::new();
        
        for (i, field) in self.schema.fields().iter().enumerate() {
            let array = batch.column(i);
            let value = self.array_to_value(array, field.data_type())?;
            map.insert(field.name().clone(), value);
        }
        
        serde_json::from_value(serde_json::Value::Object(map))
            .map_err(|e| anyhow!("Failed to deserialize struct: {}", e))
    }
    
    // 辅助方法：将JSON值转换为Arrow数组
    fn value_to_array(&self, value: &serde_json::Value, data_type: &ArrowDataType) -> Result<ArrayRef> {
        match data_type {
            ArrowDataType::Boolean => {
                let values: Vec<bool> = serde_json::from_value(value.clone())?;
                Ok(Arc::new(BooleanArray::from(values)))
            },
            ArrowDataType::Int32 => {
                let values: Vec<i32> = serde_json::from_value(value.clone())?;
                Ok(Arc::new(Int32Array::from(values)))
            },
            ArrowDataType::Int64 => {
                let values: Vec<i64> = serde_json::from_value(value.clone())?;
                Ok(Arc::new(Int64Array::from(values)))
            },
            ArrowDataType::Float32 => {
                let values: Vec<f32> = serde_json::from_value(value.clone())?;
                Ok(Arc::new(Float32Array::from(values)))
            },
            ArrowDataType::Float64 => {
                let values: Vec<f64> = serde_json::from_value(value.clone())?;
                Ok(Arc::new(Float64Array::from(values)))
            },
            ArrowDataType::Utf8 => {
                let values: Vec<String> = serde_json::from_value(value.clone())?;
                Ok(Arc::new(StringArray::from(values)))
            },
            ArrowDataType::Timestamp(_, _) => {
                let values: Vec<i64> = serde_json::from_value(value.clone())?;
                Ok(Arc::new(TimestampNanosecondArray::from(values)))
            },
            _ => Err(anyhow!("Unsupported Arrow data type: {:?}", data_type)),
        }
    }
    
    // 辅助方法：将Arrow数组转换为JSON值
    fn array_to_value(&self, array: &ArrayRef, data_type: &ArrowDataType) -> Result<serde_json::Value> {
        match data_type {
            ArrowDataType::Boolean => {
                let values = <BooleanArray as ArrowConverter<bool>>::from_arrow_array(array)?;
                Ok(serde_json::to_value(values)?)
            },
            ArrowDataType::Int32 => {
                let values = <Int32Array as ArrowConverter<i32>>::from_arrow_array(array)?;
                Ok(serde_json::to_value(values)?)
            },
            ArrowDataType::Int64 => {
                let values = <Int64Array as ArrowConverter<i64>>::from_arrow_array(array)?;
                Ok(serde_json::to_value(values)?)
            },
            ArrowDataType::Float32 => {
                let values = <Float32Array as ArrowConverter<f32>>::from_arrow_array(array)?;
                Ok(serde_json::to_value(values)?)
            },
            ArrowDataType::Float64 => {
                let values = <Float64Array as ArrowConverter<f64>>::from_arrow_array(array)?;
                Ok(serde_json::to_value(values)?)
            },
            ArrowDataType::Utf8 => {
                let values = <StringArray as ArrowConverter<String>>::from_arrow_array(array)?;
                Ok(serde_json::to_value(values)?)
            },
            ArrowDataType::Timestamp(_, _) => {
                let values = <TimestampNanosecondArray as ArrowConverter<DateTime<Utc>>>::from_arrow_array(array)?;
                Ok(serde_json::to_value(values)?)
            },
            _ => Err(anyhow!("Unsupported Arrow data type: {:?}", data_type)),
        }
    }
}

/// DataFrame 转换器
pub struct DataFrameConverter {
    serializer: StructSerializer,
}

impl DataFrameConverter {
    /// 创建新的 DataFrameConverter
    pub fn new() -> Self {
        // 定义市场数据的基本结构
        let fields = vec![
            Field::new("price", ArrowDataType::Float64, false),
            Field::new("timestamp", ArrowDataType::Timestamp(arrow::datatypes::TimeUnit::Nanosecond, None), false),
        ];
        Self {
            serializer: StructSerializer::new(fields),
        }
    }

    /// 将数据转换为 Arrow 格式
    pub fn to_arrow(&self, data: &[f64]) -> PyResult<PyObject> {
        Python::with_gil(|py| {
            // 创建价格数组
            let price_array = Float64Array::from(data.to_vec());
            
            // 创建时间戳数组（使用当前时间作为基准）
            let now = Utc::now();
            let timestamps: Vec<i64> = (0..data.len())
                .map(|i| now.timestamp_nanos_opt().unwrap_or(0) + i as i64 * 1_000_000_000) // 每个数据点间隔1秒
                .collect();
            let timestamp_array = TimestampNanosecondArray::from(timestamps.clone());

            // 创建 RecordBatch
            let batch = RecordBatch::try_new(
                Arc::new(Schema::new(vec![
                    Field::new("price", ArrowDataType::Float64, false),
                    Field::new("timestamp", ArrowDataType::Timestamp(arrow::datatypes::TimeUnit::Nanosecond, None), false),
                ])),
                vec![
                    Arc::new(price_array),
                    Arc::new(timestamp_array),
                ],
            ).map_err(|e| PyErr::new::<pyo3::exceptions::PyValueError, _>(format!("Failed to create RecordBatch: {}", e)))?;

            // 将 RecordBatch 转换为 Python 字典
            let dict = PyDict::new(py);
            
            // 添加价格数据
            let prices: Vec<f64> = data.to_vec();
            let py_prices = prices.to_object(py);
            dict.set_item("prices", py_prices)?;
            
            // 添加时间戳数据
            let py_timestamps = timestamps.to_object(py);
            dict.set_item("timestamps", py_timestamps)?;
            
            // 添加 Schema 信息
            let schema_str = batch.schema().to_string();
            dict.set_item("schema", schema_str.to_object(py))?;

            Ok(dict.to_object(py))
        })
    }
}